# -*- coding: utf-8 -*-
# @Time    : 2024/11/21 15:49
# @Author  : sjh
# @Site    : 
# @File    : filters.py
# @Comment :
import numpy as np
def low_pass_filter(data, alpha=0.1):
    """
    对3D关键点数据进行低通滤波。
    :param data: 输入的3D关键点数据，形状为 (帧数, 关键点数, 3)
    :param alpha: 滤波系数
    :return: 滤波后的3D关键点数据
    """
    filtered_data = np.zeros_like(data)
    filtered_data[0] = data[0]  # 初始化第一个帧
    for t in range(1, data.shape[0]):
        filtered_data[t] = alpha * data[t] + (1 - alpha) * filtered_data[t - 1]
    return filtered_data

class OneEuroFilter:
    def __init__(self, x0, dx0=0.0, min_cutoff=1.0, beta=0.0, d_cutoff=1.0):
        self.min_cutoff = min_cutoff
        self.beta = beta
        self.d_cutoff = d_cutoff
        self.x_prev = x0
        self.dx_prev = dx0

    def smoothing_factor(self, t_e, cutoff):
        r = 2 * np.pi * cutoff * t_e
        return r / (r + 1)

    def __call__(self, x, t_e=0.033):
        dx = (x - self.x_prev) / t_e
        dx_hat = self.smoothing_factor(t_e, self.d_cutoff) * dx + (1 - self.smoothing_factor(t_e, self.d_cutoff)) * self.dx_prev
        cutoff = self.min_cutoff + self.beta * abs(dx_hat)
        x_hat = self.smoothing_factor(t_e, cutoff) * x + (1 - self.smoothing_factor(t_e, cutoff)) * self.x_prev
        self.x_prev = x_hat
        self.dx_prev = dx_hat
        return x_hat
def moving_average_filter(data, window_size=5):
    """
    对3D关键点数据进行移动平均滤波。
    :param data: 输入的3D关键点数据，形状为 (帧数, 关键点数, 3)
    :param window_size: 滑动窗口大小
    :return: 滤波后的3D关键点数据
    """
    filtered_data = np.zeros_like(data)
    for t in range(data.shape[0]):
        start_idx = max(0, t - window_size + 1)
        filtered_data[t] = np.mean(data[start_idx:t+1], axis=0)
    return filtered_data
from scipy.signal import savgol_filter

def savitzky_golay_filter(data, window_length=5, polyorder=2):
    """
    对3D关键点数据进行Savitzky-Golay滤波。
    :param data: 输入的3D关键点数据，形状为 (帧数, 关键点数, 3)
    :param window_length: 滑动窗口大小
    :param polyorder: 多项式阶数
    :return: 滤波后的3D关键点数据
    """
    filtered_data = np.zeros_like(data)
    for joint in range(data.shape[1]):
        for dim in range(3):  # 分别对x, y, z进行滤波
            filtered_data[:, joint, dim] = savgol_filter(data[:, joint, dim], window_length, polyorder)
    return filtered_data
from scipy.signal import butter, filtfilt

def butterworth_filter(data, cutoff=0.1, fs=30, order=2):
    """
    对3D关键点数据进行Butterworth低通滤波。
    :param data: 输入的3D关键点数据，形状为 (帧数, 关键点数, 3)
    :param cutoff: 截止频率（归一化为0~0.5之间）
    :param fs: 采样频率
    :param order: 滤波器阶数
    :return: 滤波后的3D关键点数据
    """
    nyquist = 0.5 * fs  # 奈奎斯特频率
    normalized_cutoff = cutoff / nyquist
    b, a = butter(order, normalized_cutoff, btype='low', analog=False)
    filtered_data = np.zeros_like(data)
    for joint in range(data.shape[1]):
        for dim in range(3):  # 分别对x, y, z进行滤波
            filtered_data[:, joint, dim] = filtfilt(b, a, data[:, joint, dim])
    return filtered_data
if __name__ =='__main__':
    import matplotlib.pyplot as plt

    Filter3D = Filter3D()
    # 示例3D关键点数据 (帧数=50, 关键点数=28, 每点3维)
    np.random.seed(42)
    dataOri = np.cumsum(np.random.randn(50, 28, 3) * 0.1, axis=0)  # 随机生成模拟数据
    # 调用滤波方法
    lowpass_data = low_pass_filter(dataOri, alpha=0.2)
    savgol_data = savitzky_golay_filter(dataOri, window_length=5, polyorder=2)
    butterworth_data = butterworth_filter(dataOri, cutoff=0.1, fs=30, order=2)

    # 可视化某一关键点的x, y, z坐标变化（以第17个关键点为例）
    joint_index = 17
    time = np.arange(dataOri.shape[0])

    plt.figure(figsize=(12, 8))

    # 原始数据
    plt.subplot(3, 1, 1)
    plt.plot(time, dataOri[:, joint_index, 0], label="Original", alpha=0.5)
    plt.plot(time, lowpass_data[:, joint_index, 0], label="Low Pass", linestyle="--")
    plt.plot(time, savgol_data[:, joint_index, 0], label="Savitzky-Golay", linestyle="-.")
    plt.plot(time, butterworth_data[:, joint_index, 0], label="Butterworth", linestyle=":")
    plt.title(f"Joint {joint_index} - X Coordinate")
    plt.xlabel("Time")
    plt.ylabel("X Value")
    plt.legend()
    plt.grid()

    # Y 方向数据
    plt.subplot(3, 1, 2)
    plt.plot(time, dataOri[:, joint_index, 1], label="Original", alpha=0.5)
    plt.plot(time, lowpass_data[:, joint_index, 1], label="Low Pass", linestyle="--")
    plt.plot(time, savgol_data[:, joint_index, 1], label="Savitzky-Golay", linestyle="-.")
    plt.plot(time, butterworth_data[:, joint_index, 1], label="Butterworth", linestyle=":")
    plt.title(f"Joint {joint_index} - Y Coordinate")
    plt.xlabel("Time")
    plt.ylabel("Y Value")
    plt.legend()
    plt.grid()

    # Z 方向数据
    plt.subplot(3, 1, 3)
    plt.plot(time, dataOri[:, joint_index, 2], label="Original", alpha=0.5)
    plt.plot(time, lowpass_data[:, joint_index, 2], label="Low Pass", linestyle="--")
    plt.plot(time, savgol_data[:, joint_index, 2], label="Savitzky-Golay", linestyle="-.")
    plt.plot(time, butterworth_data[:, joint_index, 2], label="Butterworth", linestyle=":")
    plt.title(f"Joint {joint_index} - Z Coordinate")
    plt.xlabel("Time")
    plt.ylabel("Z Value")
    plt.legend()
    plt.grid()

    plt.tight_layout()
    plt.show()